package ssau.knyazev.fft.modules.oob;

import ssau.knyazev.core.Complex;
import ssau.knyazev.core.CComplex;
import ssau.knyazev.fft.common.FFTConst;
import ssau.knyazev.modules.Generator;
import ssau.knyazev.modules.Printer;

public class CFFTModule {

	protected CFFTModule() { }

	public static Complex[][] fft(float[][] src) {
		return fft(src, FFTConst.SIMPLE_MODE);
	}

	public static Complex[][] fft(float[][] src, int mode) {
		CComplex[][] arr = transform(src, mode);
		return fft(arr, mode);
	}

	public static Complex[][] fft(Complex[][] src) {
		return fft(src, FFTConst.SIMPLE_MODE);
	}

	public static Complex[][] fft(Complex[][] src, int mode) {
		CComplex[][] arr = transform(src, mode);
		return fft(arr, mode);
	}

	public static Complex[][] fft(CComplex[][] src) {
		return fft(src, FFTConst.SIMPLE_MODE);
	}

	public static Complex[][] fft(CComplex[][] src, int mode) {
		inverse(src);
		if (mode == FFTConst.CCOMPLEX_MIRROR_MODE) {
			return fftMirror(src);
		}
		int N = src.length;
		int size = 0;
		for (int x = src.length; x > 1; x = x >> 1){
			size++;
		}
		
		CComplex[] wi_list = new CComplex[src.length];
		CComplex[] wj_list = new CComplex[src.length];
		for (int i = 0; i < src.length; i++){
			wi_list[i] = CComplex.W_I(N, i);
			wj_list[i] = CComplex.W_J(N, i);
		}
		
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int step2 = step << 1;
			for (int down = 0; down < N - step; down++) {
				for (int right = 0; right < N - step; right++) {
					boolean flag = (down/step) % 2 + (right/step) % 2 > 0 ? false : true;
					if (flag) {
						CComplex t1 = wi_list[(down % step2)*(N/step2)];
						CComplex t2 = wj_list[(right % step2)*(N/step2)];
						butterfly(src[down][right], src[down][right + step], src[down + step][right], src[down + step][right + step], t1, t2);
					}
				}
			}
		}
		return convertToComplex(src);
	}
	
	protected static Complex[][] fftMirror(CComplex[][] src) {
		int N = src.length;
		int size = 0;
		for (int x = src.length; x > 1; x = x >> 1){
			size++;
		}
		
		CComplex[] wi_list = new CComplex[src.length];
		CComplex[] wj_list = new CComplex[src.length];
		for (int i = 0; i < src.length; i++){
			wi_list[i] = CComplex.W_I(N, i);
			wj_list[i] = CComplex.W_J(N, i);
		}
		
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int matrixWidth = step << 1;
			int part = matrixWidth/4;
			for (int down = 0; down < N - step; down++) {
				for (int right = 0; right < N - step; right++) {
					int dmatrInd = down/matrixWidth * matrixWidth;
					int rmatrInd = right/matrixWidth * matrixWidth;
					int tdown = down - dmatrInd;
					int tright = right - rmatrInd;
					boolean flag = (tdown <= part && tright <= part);
					if (flag) {
						CComplex t1 = wi_list[(down % matrixWidth)*(N/matrixWidth)];
						CComplex t2 = wj_list[(right % matrixWidth)*(N/matrixWidth)];
						butterfly(src[down][right], src[down][right + step], src[down + step][right], src[down + step][right + step], t1, t2);
						
						if (q >= 2){
							int rmirr = mirrorIndex(right, matrixWidth);//(rmatrInd == right) ? right : matrixWidth - right + 2*rmatrInd;
							int dmirr = mirrorIndex(down, matrixWidth); // (dmatrInd == down) ? down : matrixWidth - down + 2*dmatrInd;
							
							mirror(src, down, right, dmirr, rmirr);
							
							rmirr = mirrorIndex(right + step, matrixWidth);
							mirror(src, down, right + step, dmirr, rmirr);
							
							dmirr = mirrorIndex(down + step, matrixWidth);
							rmirr = mirrorIndex(right, matrixWidth);
							mirror(src, down + step, right, dmirr, rmirr);
							
							rmirr = mirrorIndex(right + step, matrixWidth);
							mirror(src, down + step, right + step, dmirr, rmirr);
						}

					}
				}
			}
		}
		
		return null;// convertToComplex(src);
	}

	protected static CComplex[][] transform(float[][] in, int mode) {
		CComplex[][] res = null;
		if (mode == FFTConst.CCOMPLEX_HALF_MODE) {
			res = new CComplex[in.length / 2][in.length / 2];
			for (int i = 0; i < in.length / 2; i++){
				for (int j = 0; j < in.length / 2; j++) {
					res[i][j] = new CComplex(in[2 * i][2 * j],
							in[2 * i][2 * j + 1], in[2 * i + 1][2 * j],
							in[2 * i + 1][2 * j + 1]);
				}
			}
		} else {
			res = new CComplex[in.length][in.length];
			for (int i = 0; i < in.length; i++){
				for (int j = 0; j < in.length; j++) {
					res[i][j] = new CComplex(in[i][j], 0, 0, 0);
				}
			}
		}
		return res;
	}

	protected static CComplex[][] transform(Complex[][] in, int mode) {
		if (mode == FFTConst.CCOMPLEX_HALF_MODE) {
			CComplex[][] res = new CComplex[in.length / 2][in.length / 2];
			for (int i = 0; i < in.length / 2; i++)
				for (int j = 0; j < in.length / 2; j++) {
					res[i][j] = new CComplex(in[2 * i][2 * j].getReal(),
							in[2 * i][2 * j + 1].getReal(),
							in[2 * i + 1][2 * j].getReal(),
							in[2 * i + 1][2 * j + 1].getReal());
				}
			return res;
		} else {
			CComplex[][] res = new CComplex[in.length][in.length];
			for (int i = 0; i < in.length; i++)
				for (int j = 0; j < in.length; j++) {
					res[i][j] = new CComplex(in[i][j].getReal(), 0, 0, 0);
				}
			return res;
		}
	}

	protected static void butterfly(CComplex a00, CComplex a01, CComplex a10, CComplex a11, CComplex wi, CComplex wj) {
		CComplex newa00 = a00.clone();
		CComplex newa10 = a00.clone();
		CComplex newa01 = a00.clone();
		CComplex newa11 = a00.clone();
		
		CComplex t = wi.multy(a10);
		newa00 = newa00.add(t);
		newa10 = newa10.sub(t);
		newa01 = newa01.add(t);
		newa11 = newa11.sub(t);
		
		t = wj.multy(a01);
		newa00 = newa00.add(t);
		newa10 = newa10.add(t);
		newa01 = newa01.sub(t);
		newa11 = newa11.sub(t);
		
		t = wj.multy(a11);
		t = wi.multy(t);
		newa00 = newa00.add(t);
		newa10 = newa10.sub(t);
		newa01 = newa01.sub(t);
		newa11 = newa11.add(t);

		copyValues(a00, newa00);
		copyValues(a01, newa01);
		copyValues(a10, newa10);
		copyValues(a11, newa11);
	}

	protected static void copyValues(CComplex a, CComplex b) {
		a.setA(b.getA());
		a.setB(b.getB());
		a.setC(b.getC());
		a.setD(b.getD());
	}
	
	protected static int mirrorIndex(int srcIndex, int matrixWidth){
		int matrixIndex = srcIndex/matrixWidth * matrixWidth;
		return (matrixIndex == srcIndex) ? srcIndex : matrixWidth - srcIndex + 2*matrixIndex;
	}
	
	protected static void mirror(CComplex[][] src, int down, int right, int dmirr, int rmirr){
		if (rmirr != right){
			src[down][rmirr].setA(src[down][right].getA());
			src[down][rmirr].setB(src[down][right].getB());
			src[down][rmirr].setC(-src[down][right].getC());
			src[down][rmirr].setD(-src[down][right].getD());
		}
		if (dmirr != down){
			src[dmirr][right].setA(src[down][right].getA());
			src[dmirr][right].setB(-src[down][right].getB());
			src[dmirr][right].setC(src[down][right].getC());
			src[dmirr][right].setD(-src[down][right].getD());
		}
		if (rmirr != right && dmirr != down){
			src[dmirr][rmirr].setA(src[down][right].getA());
			src[dmirr][rmirr].setB(-src[down][right].getB());
			src[dmirr][rmirr].setC(-src[down][right].getC());
			src[dmirr][rmirr].setD(src[down][right].getD());
		}
	}

	protected static Complex[][] convertToComplex(CComplex[][] arr) {
		Complex[][] res = new Complex[arr.length][arr.length];
		for (int i = 0; i < res.length; i++){
			for (int j = 0; j < res.length; j++){
				res[i][j] = arr[i][j].getComplex();
			}
		}
		return res;
	}

	protected static void inverse(CComplex[] arr) {
		int dl = arr.length / 2;
		int st = arr.length - 1;
		int j = 0;
		for (int i = 0; i < st; i++) {
			if (i < j) {
				CComplex s0 = arr[i];
				arr[i] = arr[j];
				arr[j] = s0;
			}
			int k = dl;
			while (k <= j) {
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}

	protected static void inverse(CComplex[][] src) {
		int N = src.length;
		for (int i = 0; i < N; i++) {
			inverse(src[i]);
		}
		CComplex[] arr = new CComplex[N];
		for (int i = 0; i < N; i++) {
			for (int j = 0; j < N; j++) {
				arr[j] = src[j][i];
			}
			inverse(arr);
			for (int j = 0; j < N; j++) {
				src[j][i] = arr[j];
			}
		}
	}
	
	public static void main(String[] args){
		float[][] arr = Generator.generateSimpleMatrix(4096);
		//Printer.printVector(arr);
		long sum = 0;
		for (int i = 0; i < 10; i++){
			long time = System.nanoTime();
			Complex[][] res = CFFTModule.fft(arr, FFTConst.CCOMPLEX_MIRROR_MODE);
			time = System.nanoTime() - time;
			sum += time;
		}	
		System.out.println("RESULT");
		//Printer.printVector(res);
			
		sum /= 10;
		System.out.println("\n" + sum);
	}

}
